#!/usr/bin/env python

#$Id: flatten.py 15 2006-04-11 23:22:57Z magee $
# -------------------------------------------------------------
#
__version__      = '$Revision: 15 $ '[11:-3]
__version_date__ = '$Date: 2006-04-11 23:22:57 +0000 (Tue, 11 Apr 2006) $ '[7:-3]
__author__       = 'R. Bouwens, <bouwens@ucolick.org>, D. Magee, <magee@ucolick.org>'

import numarray as N
import numarray.nd_image as nd
from numarray import ieeespecial as ieee
import pyfits,os,imagestats,glob,sys
from pysqlite2 import dbapi2 as sqlite
import shutil

def measurenoise(flatdata,dq,staticmask,clip):
    ndq = (dq == 0) * (staticmask == 1) * (flatdata < clip) * (flatdata > -clip*1.3)
    hclip = clip / 2
    flatdata_med = nd.median_filter(flatdata,size=3)
    ndq *= (flatdata_med < hclip)
    dqd = (ndq == 1)
    good = flatdata[dqd]
    stats = imagestats.ImageStats(good, fields='stddev')
    return stats.stddev

def findbkg(data,flatdata,dq,staticmask,clip,size,goodfact=0.25):
    ndq = (dq == 0) * (staticmask == 1) * (flatdata < clip) * (flatdata > -clip*1.3)
    hclip = clip / 2
    flatdata_med = nd.median_filter(flatdata,size=3)
    ndq *= (flatdata_med < hclip)

    (xdim,ydim) = data.shape
    ngridx = xdim / size
    ngridy = ydim / size

    Median = N.zeros((ngridx,ngridy),'Float64')
    NContr = N.zeros((ngridx,ngridy),'Int32')
    for i in range(ngridx):
        for j in range(ngridy):
            xl = i*size
            xh = N.clip(xl+size,0,xdim)
            yl = j*size
            yh = N.clip(yl+size,0,ydim)
            datasl = data[xl:xh,yl:yh]
            dqsl = ndq[xl:xh,yl:yh]
            good = (dqsl == 1)
            curdata = datasl[good]
            NContr[i,j] = len(curdata)
            if NContr[i,j] > 2:
                stats = imagestats.ImageStats(curdata, fields='median',binwidth=0.00005)
                Median[i,j] = stats.median
            else:
                Median[i,j] = 0
    MedianM = N.zeros((ngridx+1,ngridy+1),'Float64')
    for i in range(ngridx):
        for j in range(ngridy):
            notdone = 1
            sizem = 1
            while notdone:
                xl = N.clip(i-sizem,0,ngridx)
                xh = N.clip(i+sizem+1,0,ngridx)
                yl = N.clip(j-sizem,0,ngridy)
                yh = N.clip(j+sizem+1,0,ngridy)  
                datasl = Median[xl:xh,yl:yh]
                dqsl = NContr[xl:xh,yl:yh]
                good = (dqsl > size*size*goodfact)
                curdata = datasl[good]
                if len(curdata) > 1:
                    stats = imagestats.ImageStats(curdata, fields='median',binwidth=0.00005)
                    if len(ieee.getnan(N.array([stats.median]))[0]) == 0:
                        notdone = 0
                if notdone:
                    sizem += 1

            MedianM[i+1,j+1] = stats.median
            if i == 0:
                MedianM[0,j+1] = stats.median
            if j == 0:
                MedianM[i+1,0] = stats.median
                if i == 0:
                    MedianM[0,0] = stats.median

    newmedian = nd.affine_transform(MedianM,[[1./size,0],[0,1./size]],output_shape=(256,256),offset=(0.5,0.5))
    newmedian *= staticmask
    return newmedian

def findbkgquad(data,flatdata,dq,staticmask,clip,size,goodfact=0.5):
  bkg = N.zeros((256,256),type=N.Float32)
  for m in range(2):
    for n in range(2):
      xl = m*128
      xh = m*128+128
      yl = n*128
      yh = n*128+128
      cbkg = findbkg(data[yl:yh,xl:xh],data[yl:yh,xl:xh],dq[yl:yh,xl:xh],staticmask[yl:yh,xl:xh],0.01,size=9)
      bkg[yl:yh,xl:xh] = cbkg
  return bkg

def flatten(inname, outname, staticmask):
    newf = pyfits.open(inname)
    data = newf['SCI'].data
    dq = pyfits.getdata(inname,ext='DQ')
    staticmask = pyfits.getdata(staticmask)
    cmedian = findbkg(data,data,dq,staticmask,0.08,size=11)
    flatdata = data - cmedian
    noise = measurenoise(flatdata,dq,staticmask,0.03)
    print 'RMS noise = ',noise
    cmedian = findbkg(data,flatdata,dq,staticmask,noise*2,size=11)
    data -= cmedian
    newf.writeto(outname,clobber=True)
    newf.close()
